import copy
import struct

import numpy as np
from PySide2.QtGui import QPixmap, qRgb, QImage
from PIL import Image


# 结构元素
MORPH_CROSS = np.array([[0, 1, 0],
                        [1, 1, 1],
                        [0, 1, 0]], int)

MORPH_RECT = np.array([[1, 1, 1],
                       [1, 1, 1],
                       [1, 1, 1]], int)


def open_image(path):
    dst = np.zeros((400, 300), np.uint8)
    if path.endswith('.raw'):
        dst = read_raw(path)
    elif path.endswith('.bmp') or path.endswith('.BMP'):
        dst = read_bmp(path)
    return dst


# def read_raw(path):
#     w, h = np.fromfile(path, dtype=np.uint32)[:2]  # .raw格式图片前两位表示图像的宽和高，这个宽和高一般大于256，且占各4个字节（32bit），所以先用np.uint32去读宽和高
#     dst = np.fromfile(path, dtype=np.uint8)[8:]  # 排除前8个字节，其余每个字节均为一个像素的灰度级，用np.uint8
#     dst = dst.reshape(h, w)  # 读出来的值为一维数组，重塑形为h x w 的二维矩阵
#     return dst


def read_raw(path):
    with open(path, 'rb') as f:
        data = f.read()
    w_byte = data[:4]
    h_byte = data[4:8]
    im = data[8:]
    w = struct.unpack('I', w_byte)[0]
    h = struct.unpack('I', h_byte)[0]
    dst = np.zeros((h, w), dtype=np.uint8)

    for i in range(h):
        for j in range(w):
            dst[i][j] = im[i*w + j]
    return dst

# 读取bmp格式，关于位图的读取，可参考https://blog.csdn.net/qq_43409114/article/details/104538619，这里直接用PIL了
def read_bmp(path):
    im = Image.open(path)
    dst = np.array(im)
    return dst


# 显示图像，image为图像像素灰度的矩阵，label为用于装载图片的组件，x, y为label在界面的位置
def show_image(image, label, x, y):
    h, w = image.shape[:2]
    qimg = QImage(w, h, QImage.Format_RGB32)
    # 彩色图像
    if image.ndim == 3:
        for i in range(h):
            for j in range(w):
                val = qRgb(image[i][j][0], image[i][j][1], image[i][j][2])  # rgb
                qimg.setPixel(j, i, val)
    # 灰度图像
    elif image.ndim == 2:
        for i in range(h):
            for j in range(w):
                val = qRgb(image[i][j], image[i][j], image[i][j])
                qimg.setPixel(j, i, val)
    pixmap = QPixmap.fromImage(qimg)
    label.setPixmap(pixmap)
    label.setGeometry(x, y, w, h)
    label.show()
    label.raise_()  # 置于最上层


# 计算直方图和生成直方图的图像
def cal_hist(src):
    hist = np.zeros(256, dtype=int)
    # TODO: 计算直方图

    #
    return hist

# 绘制直方图的图像
def plot_hist(src):
    hist = cal_hist(src)
    # 直方图单独看作一张图片，设置其高度为400px，宽度为1024px，即为256的4倍，这样能够显示出一条条分立的谱线
    img_h = 400
    img_w = 256 * 4
    max_val = hist.max()
    if max_val == 0:
        max_val = 1000000
    coeffs = (400 - 20) / max_val  # 直方图归一化，减20是为了不到顶
    normed_hist = hist * coeffs
    normed_hist = normed_hist.astype(int)

    dst = np.ones((img_h, img_w), dtype=np.uint8) + 254  # 用于保存直方图的图片的矩阵，全白
    for i in range(img_w >> 2):  # 右移两位，相当于除以4
        for j in range(img_h, 0, -1):
            if j < normed_hist[i]:
                dst[400 - j][i << 2] = 0
    return dst


# 二值化
def binarize(src, th):
    h, w = src.shape[:2]
    dst = np.zeros_like(src)
    # TODO: 二值化

    return dst

#
# 幂次变换
def exp_trans(src, expon):
    dst = np.zeros_like(src)
    # TODO: 幂次变换

    return dst



# 幂次变换-3通道
def exp_trans_bmp(src, expon):
    h, w = src.shape[:2]
    src_r = src[:, :, 0]
    src_g = src[:, :, 1]
    src_b = src[:, :, 2]
    dst_r = exp_trans(src_r, expon)
    dst_g = exp_trans(src_g, expon)
    dst_b = exp_trans(src_b, expon)
    dst = np.zeros((h, w, 3))
    dst[:, :, 0] = dst_r
    dst[:, :, 1] = dst_g
    dst[:, :, 2] = dst_b
    return dst


# 对数变换
def log_trans(src):
    dst = np.zeros_like(src)
    # TODO: 对数变换

    return dst


# 对数变换-3通道
def log_trans_bmp(src):
    h, w = src.shape[:2]
    src_r = src[:, :, 0]
    src_g = src[:, :, 1]
    src_b = src[:, :, 2]
    dst_r = log_trans(src_r)
    dst_g = log_trans(src_g)
    dst_b = log_trans(src_b)
    dst = np.zeros((h, w, 3))
    dst[:, :, 0] = dst_r
    dst[:, :, 1] = dst_g
    dst[:, :, 2] = dst_b
    return dst


# 直方图均衡
def hist_eql(src):
    dst = np.zeros_like(src)
    # TODO: 直方图均衡

    return dst


# sobel算子边缘提取
def sobel(src):
    sobel_x = np.zeros_like(src)
    sobel_y = np.zeros_like(src)
    # TODO: Sobel算子边缘提取

    return sobel_x, sobel_y


# 腐蚀
def erode(src, iterations=1):
    dst = copy.deepcopy(src)
    # TODO: 腐蚀

    return dst


# 膨胀
def dilate(src, iterations=1):
    dst = copy.deepcopy(src)
    # TODO: 膨胀

    return dst


# 开运算
def opening(src, iterations):
    dst = copy.deepcopy(src)
    # TODO: 开运算
    return dst


# 闭运算
def closing(src, iterations):
    dst = copy.deepcopy(src)
    # TODO: 闭运算

    return dst


# 均值滤波
def mean_filter(src):
    h, w = src.shape[:2]
    dst = np.zeros((h, w), dtype=np.uint8)
    # TODO: 均值滤波

    return dst


# 中值滤波
def median_filter(src):
    h, w = src.shape[:2]
    dst = np.zeros((h, w), dtype=np.uint8)
    # TODO: 中值滤波

    return dst


# 高斯滤波
def gaussian_filter(src):
    h, w = src.shape[:2]
    dst = np.zeros((h, w), dtype=np.uint8)
    # TODO: 高斯滤波

    return dst


# 连通域分析
def conc(src):
    h, w = src.shape[:2]
    labels = np.zeros((h, w))
    # TODO: 连通域分析

    return labels


# 平移变换
def translate(src, delta_x, delta_y):
    h, w = src.shape[:2]
    # 若平移的偏移量大于图像的宽高，则取其对应的模宽或高的运算结果
    if delta_x >= w or delta_x <= -w:
        delta_x = delta_x % w
    if delta_y >= h or delta_y <= -h:
        delta_y = delta_y % h

    dst = np.zeros_like(src)
    M = np.array([[1, 0, delta_x],
                  [0, 1, delta_y]])
    for i in range(h):
        for j in range(w):
            position = np.array([[j],
                                 [i],
                                 [1]])

            new_position = M @ position  # 矩阵相乘，得到变换后的位置
            x, y = new_position[:, 0]
            # 移除图像外的点不管
            if x >= w:
                continue
                # x = x % w
            if y >= h:
                continue
                # y = y % h
            dst[y][x] = src[i][j]
    return dst


# 翻转变换
def flip(src, M):
    h, w = src.shape[:2]
    dst = np.zeros_like(src)
    for i in range(h):
        for j in range(w):
            position = np.array([[j],
                                 [i],
                                 [1]])

            new_position = M @ position  # 矩阵相乘，得到变换后的位置
            x, y = new_position[:, 0]
            dst[y][x] = src[i][j]
    return dst


# 缩放
def zoom(src, f_x, f_y):
    h, w = src.shape[:2]
    new_h = int(h * f_y)
    new_w = int(w * f_x)
    dst = np.zeros((new_h, new_w), dtype=np.uint8)

    M = np.array([[f_x, 0, 0],
                  [0, f_y, 0]])
    for i in range(h):
        for j in range(w):
            position = np.array([[j],
                                 [i],
                                 [1]])

            new_position = M @ position
            x, y = new_position[:, 0]
            x = int(x)
            y = int(y)
            dst[y][x] = src[i][j]
    tmp = np.zeros((new_h + 2, new_w + 2), dtype=np.uint8)
    tmp[1:-1, 1:-1] = copy.deepcopy(dst)
    tmp[:, -1] = tmp[:, 1]
    tmp[:, 0] = tmp[:, -2]
    tmp[-1, :] = tmp[1, :]
    tmp[0, :] = tmp[-2, :]
    for i in range(1, new_h + 1):
        for j in range(1, new_w + 1):
            if tmp[i][j] == 0:
                block = tmp[i - 1:i + 2, j - 1:j + 2]  # 取(i,j)为中心的3x3的一个小块
                new_val = mean_value(block)
                dst[i - 1][j - 1] = new_val
    return dst


# 旋转
def rotate(src, theta):
    theta = theta / 180 * np.pi  # 角度转弧度

    h, w = src.shape[:2]
    cx, cy = w >> 1, h >> 1  # 图片中心
    M = np.array([[+np.cos(theta), +np.sin(theta), (1 - np.cos(theta) * cx - cy * np.sin(theta))],
                  [-np.sin(theta), +np.cos(theta), (1 - np.cos(theta) * cy + cx * np.sin(theta))]])
    new_h = int(w * np.sin(theta) + h * np.cos(theta))
    new_w = int(h * np.sin(theta) + w * np.cos(theta))
    dst = np.zeros((new_h, new_w), dtype=np.uint8)
    new_cx, new_cy = new_w >> 1, new_h >> 1

    for i in range(h):
        for j in range(w):
            position = np.array([[j],
                                 [i],
                                 [1]])

            new_position = M @ position
            x, y = new_position[:, 0]
            x = int(x + new_cx)
            y = int(y + new_cy)  # 移到中心
            if x < 0:
                continue
            if y < 0:
                continue
            if x >= new_w:
                x -= new_w
            if y >= new_h:
                y -= new_h
            dst[y][x] = src[i][j]
    # 当倾斜角度不是90的整数倍时会产生缝隙，用一个像素周围非零像素的均值来填充这些缝隙
    tmp = np.zeros((new_h+2, new_w+2), dtype=np.uint8)
    tmp[1:-1, 1:-1] = copy.deepcopy(dst)
    tmp[:, -1] = tmp[:, 1]
    tmp[:, 0] = tmp[:, -2]
    tmp[-1, :] = tmp[1, :]
    tmp[0, :] = tmp[-2, :]
    for i in range(1, new_h+1):
        for j in range(1, new_w+1):
            if tmp[i][j] == 0:
                block = tmp[i-1:i+2, j-1:j+2]  # 取(i,j)为中心的3x3的一个小块
                new_val = mean_value(block)
                dst[i-1][j-1] = new_val
    return dst


# 3x3非零均值填充
def mean_value(block):
    seq = block.reshape(-1)
    count = 0
    sigma = 0
    for i in range(9):
        if seq[i] != 0:
            count += 1
            sigma += seq[i]
    if count == 0:
        count = 1
    return int(sigma / count)
